﻿#include "libmesh/linear_solver.h"
#include "libmesh/equation_systems.h"
#include "libmesh/system_subset.h"

#include "EbESystem.h"

//

void EbESystem::solve()
{
    //LinearImplicitSystem::solve();
    //ebe_solve_cg();
    ebe_solve_bicgstab();
}

void EbESystem::ebe_solve_cg()
{
   

    this->assemble();

    auto& y = add_vector("y", false);
    auto& x = add_vector("x", false);
    auto& p = add_vector("p", false);
    auto& r = add_vector("r", false);

    //x0 = 0;
    zero(x);
    // y = Ax0
    linear_operator(y, x);
    // r0 = b-y
    axpby(r, *rhs, y);
    // p0=r0
    copy(p, r);

    for (size_t k = 0; k < 200; k++)
    {
        Real rz_old = dot(r, r);
        std::cout << "iteration: " << k << " res: " << rz_old << std::endl;
        // y = Ap;
        linear_operator(y, p);
        Real alpha = dot(r, r) / dot(p, y);
        std::cout << alpha << std::endl;
        // x = x + alpha*p
        axpy(x, p, alpha);
        // r = r-alpha*y;
        axpy(r, y, -alpha);

        if (dot(r, r) < 1E-12)
        {
            std::cout << "Convergence: " << dot(r, r) << std::endl;
            break;
        }

        Real beta = dot(r, r) / rz_old;
        axpby(p, r, p, beta);
    }
    copy(*solution, x);

}

void EbESystem::ebe_solve_bicgstab()
{
    this->assemble();

    auto& y = add_vector("y", false);
    auto& x = add_vector("x", false);
    auto& p = add_vector("p", false);
    auto& r = add_vector("r", false);
    auto& s = add_vector("s", false);
    auto& As = add_vector("As", false);
    auto& r_t = add_vector("r_t", false);

    //x0 = 0;
    zero(x);
    // y = A*x0
    linear_operator(y, x);
    // r0 = b-y
    axpby(r, *rhs, y);
    // p0=r0
    copy(p, r);
    copy(r_t, r);

    for (size_t k = 0; k < 200; k++)
    {
        Real rz_old = dot(r, r_t);
        std::cout << "iteration: " << k << " res: " << rz_old << std::endl;
        // y = Ap;
        linear_operator(y, p);
        Real alpha = dot(r, r_t) / dot(y, r_t);

        for (size_t i = 0; i < s.size(); i++)
            s.set(i, r(i) - alpha * y(i));

        // As = A*s;
        linear_operator(As, s);

        Real omega = dot(As, s) / dot(As, As);

        // x = x +alpha*p+w*s
        for (size_t i = 0; i < x.size(); i++)
            x.set(i, x(i) + alpha * p(i)+omega*s(i));

        // r = s - omega*As;
        for (size_t i = 0; i < r.size(); i++)
            r.set(i, s(i) - omega * As(i));

        if (dot(r, r) < 1E-12)
        {
            std::cout << "Convergence: " << dot(r, r) << std::endl;
            break;
        }

        Real beta = alpha/omega *dot(r, r_t)/ rz_old;

        // p = r+beta*(p-omega*Ap)
        for (size_t i = 0; i < p.size(); i++)
            p.set(i, r(i) + beta* (p(i)-omega*y(i)));
    }

    copy(*solution, x);
}

//void EbESystem::ebe_solve()
//{
    //this->assemble();
       //
       //auto& diag = add_vector("diag", false);
       //matrix->get_diagonal(diag);

       //auto& y = add_vector("y", false);
       //auto& x = add_vector("x", false);
       //auto& p = add_vector("p", false);
       //auto& r = add_vector("r", false);
       //auto& z = add_vector("z", false);

       ////x0 = 0;
       //zero(x);
       //// y = Ax0
       //linear_operator(y, x);
       //// r0 = b-y
       //axpby(r, *rhs, y);
       //// p0=r0
       ////preconditioner
       //precondition(z, r);
       //copy(p, z);

       //for (size_t k = 0; k < 200; k++)
       //{
       //    Real rz_old = dot(r, z);
       //    //std::cout << "iteration: " << k << " res: " << rz_old << std::endl;
       //    // y = Ap;
       //    linear_operator(y, p);
       //    Real alpha = dot(r, z) / dot(p, y);
       //    //std::cout << alpha << std::endl;
       //    // x = x + alpha*p
       //    axpy(x, p, alpha);
       //    // r = r-alpha*y;
       //    axpy(r, y, -alpha);

       //    if (dot(r, z) < 1E-12)
       //    {
       //        std::cout << "Convergence: " << dot(r, z) << std::endl;
       //        break;
       //    }

       //    // preconditioner
       //    precondition(z, r);

       //    Real beta = dot(r, z) / rz_old;
       //    //std::cout << beta << std::endl;

       //    axpby(p, z, p, beta);
       //}
       //copy(*solution, x);
//}